import numpy as np
import cv2

def binarize_img(img, thresh=1):
    """
    将计数图二值化, 大于等于thresh的像素设为1, 其余为0
    """
    return (img >= thresh).astype(np.uint8)

def dilate_img(img, kernel_size=5, shape='rect'):
    """
    对二值图像进行膨胀操作
    :param img: 输入二值图像
    :param kernel_size: 卷积核大小
    :param shape: 'rect' 或 'ellipse'
    """
    if shape == 'ellipse':
        kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (kernel_size, kernel_size))
    else:
        kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (kernel_size, kernel_size))
    return cv2.dilate(img, kernel)

def create_eyelid_glint_mask(pos_img, neg_img, 
                             bin_thresh=1, dilate_kernel=5, dilate_shape='rect', 
                             expand_kernel=7, expand_shape='ellipse'):
    """
    根据正负极性图像生成眼睑和光斑掩码
    :param pos_img: 已去噪的正极性计数图
    :param neg_img: 已去噪的负极性计数图
    :param bin_thresh: 二值化阈值
    :param dilate_kernel: 膨胀核大小
    :param dilate_shape: 膨胀核形状
    :param expand_kernel: 最后一次膨胀核大小
    :param expand_shape: 最后一次膨胀核形状
    :return: eyelid_glint_mask
    """
    # 1. 二值化
    pos_bin = binarize_img(pos_img, bin_thresh)
    neg_bin = binarize_img(neg_img, bin_thresh)
    # 2. 膨胀
    pos_dilated = dilate_img(pos_bin, kernel_size=dilate_kernel, shape=dilate_shape)
    neg_dilated = dilate_img(neg_bin, kernel_size=dilate_kernel, shape=dilate_shape)
    # 3. 交集
    mask_temp = cv2.bitwise_and(pos_dilated, neg_dilated)
    # 4. 再膨胀
    eyelid_glint_mask = dilate_img(mask_temp, kernel_size=expand_kernel, shape=expand_shape)
    return eyelid_glint_mask
